import numpy as np
import copy
class GroupSubspace():
    
    def __init__(self, N_group, N_agent, thread_index):
        self.N_group = N_group   # each dimension for a group
        self.N_agent = N_agent
        self.agents_location = np.zeros(shape = (self.N_agent,self.N_group))
        self.group_assign_hysteresis = np.random.randint(low=0, high=self.N_group, size=self.N_agent, dtype=np.int)# random assign at init
        # P_ stand for Parameter vars
        self.P_StepRange = 0.85 # step range on chosen dimension
        self.P_PullRange = 0.5
        self.P_HysteresisLimit = 1.0
        self.internal_reward = None
        self.const0 = np.zeros(shape = (self.N_group))
        self.constM = np.ones(shape = (self.N_group)) * self.P_HysteresisLimit*3
        self.constPB = np.ones(shape = (self.N_agent,self.N_group)) * self.P_PullRange
        self.thread_index = thread_index

    def step_efficient(self, locations, group_intention_actions):
        ## locations  # dim1 agent, dim2 groups
        intention_group_mask = np.eye(self.N_group)[group_intention_actions]
        locations[intention_group_mask.astype(bool)] += self.P_StepRange
        pull_shift_1 = locations - self.constPB
        locations = np.maximum(pull_shift_1, self.const0)  # don't go <0!
        locations = np.minimum(locations, self.constM)  # don't go >M!
        return locations

    def step(self, group_intention_action):
        # assert len(group_intention_action) == self.N_agent
        # move the agent in inherent group space
        self.agents_location = self.step_efficient(self.agents_location, group_intention_action)
        self.update_hysteresis_group()
        self.update_crowding_reward_signal()

        # make obs for each agent
        agents_gstate_obs, agents_gdiv_obs = self.make_obs()
        return agents_gstate_obs, agents_gdiv_obs, self.internal_reward.copy()

    def update_crowding_reward_signal(self):
        self.internal_reward = np.zeros(shape=(self.N_agent))
        self.group_agent_cnt = np.zeros(shape=(self.N_group))
        # penelty for empty group
        # get the number of empty groups
        g = self.group_assign_hysteresis
        for x in range(self.N_agent):
            self.group_agent_cnt[g[x]] += 1
        for y in range(self.N_group):
            if self.group_agent_cnt[y] <= 0:
                self.internal_reward -= 1   # 如果出现空组别，施加惩罚 1
        pass


    def update_hysteresis_group(self):

        locations = self.agents_location
        locations_argmax = np.argmax(locations, axis=1)
        locations_update = (np.max(locations, axis=1) >= self.P_HysteresisLimit)
        self.group_assign_hysteresis = np.where(locations_update, locations_argmax, self.group_assign_hysteresis)



    def reset(self):
        self.group_assign_hysteresis = np.random.randint(low=0, high=self.N_group, size=self.N_agent, dtype=np.int)# random assign at init
        self.agents_location = np.zeros(shape = (self.N_agent,self.N_group))
        for i in range(self.N_agent):
            g = self.group_assign_hysteresis[i] # group belonging
            self.agents_location[i,g] = 1.0001  # [0,0,0,0,0,1,0,0,0]
        pass
        # make obs for each agent
        agents_gstate_obs, agents_gdiv_obs = self.make_obs()
        return agents_gstate_obs, agents_gdiv_obs


    def make_obs(self):
        '''self.agents_location.shape = (智能体数量, self.N_group) # agents_location.shape = (dim1 = nagent, dim2 = ngroups)
         neat trick of dimension
         形状是（智能体数，观测维），
         观测维前self.N_group列是自己的，后面所有列是其他智能体的（自己的观测值会再次重复一遍）'''
        tmp = self.agents_location.reshape(1, self.N_agent*self.N_group).repeat(self.N_agent,axis = 0)
        agents_gstate_obs = np.concatenate((self.agents_location, tmp), axis=1)
        # 形状是（智能体数，1+智能体数量），
        # 第二维度是 自己的组别+其他所有智能体的组别# group_assign_hysteresis.shape = (dim1 = nagent)
        vec_t = self.group_assign_hysteresis.reshape(self.N_agent, 1)  
        mat_t = self.group_assign_hysteresis.reshape(1, self.N_agent).repeat(self.N_agent,axis = 0)
        agents_gdiv_obs = np.concatenate((vec_t, mat_t),axis=1)
        return agents_gstate_obs, agents_gdiv_obs